import paho.mqtt.client as mqtt
import json
import base64
import cv2
import numpy as np
import time
import os
from datetime import datetime, timedelta
import pytz
from ultralytics import YOLO
import threading
import glob
from flask import Flask, render_template, jsonify, request, send_from_directory
import sqlite3

# 导入配置
from config import MQTT_CONFIG, FLASK_CONFIG, DATABASE_CONFIG

# 导入工具类
from utils.db_utils import get_db_connection, limit_image_files
from utils.image_utils import encode_image_to_base64

# MQTT 配置
broker = MQTT_CONFIG['broker']
port = MQTT_CONFIG['port']
feed_topic = MQTT_CONFIG['feed_topic']
device_topic = MQTT_CONFIG['device_topic']
username = MQTT_CONFIG['username']
password = MQTT_CONFIG['password']

# 创建保存目录
os.makedirs("screenshots", exist_ok=True)
os.makedirs("camera_shots", exist_ok=True)
os.makedirs("detections", exist_ok=True)
os.makedirs("templates", exist_ok=True)
os.makedirs("static", exist_ok=True)

# 存储设备状态，增加更多信息
devices = {}

# 存储历史设备信息
history_devices = {}

# 加载人形检测模型
try:
    pose_model = YOLO('models/yolov8n-pose.pt')
    print("✅ 人形检测模型已加载")
except Exception as e:
    print(f"❌ 人形检测模型加载失败: {e}")
    pose_model = None

# 人体检测日志记录器
from human_detection_logger import detection_logger

# AI分析器
from ai_analyzer import ai_analyzer

# 存储正在进行的人体检测记录
active_detections = {}  # device_id -> log_id 映射

# 定时任务线程
daily_summary_thread = None
daily_summary_running = True

# Flask 应用
app = Flask(__name__, 
            template_folder='templates',
            static_folder='static')

# 设备心跳检查线程
def check_device_heartbeat():
    """定期检查设备心跳，清理离线设备"""
    while True:
        try:
            current_time = time.time()
            offline_devices = []
            
            # 检查超过5分钟没有心跳的设备，标记为离线
            for device_id, device_info in devices.items():
                last_heartbeat = device_info.get('last_heartbeat', 0)
                if current_time - last_heartbeat > 300:  # 5分钟
                    # 标记为离线但保留在设备列表中
                    devices[device_id]['status'] = 'offline'
                    print(f"🔚 设备 {device_id} 已离线")
            
            time.sleep(60)  # 每分钟检查一次
        except Exception as e:
            print(f"❌ 设备心跳检查出错: {e}")
            time.sleep(60)

# 启动设备心跳检查线程
heartbeat_thread = threading.Thread(target=check_device_heartbeat, daemon=True)
heartbeat_thread.start()

# MQTT客户端
client = mqtt.Client()

def on_connect(client, userdata, flags, rc):
    """MQTT连接回调"""
    if rc == 0:
        print("✅ 已连接到MQTT服务器")
        # 订阅设备主题和feed主题
        client.subscribe(device_topic)
        client.subscribe(f"{feed_topic}/#")  # 订阅所有设备的feed消息
        print(f"👂 已订阅设备主题: {device_topic}")
        print(f"👂 已订阅feed主题: {feed_topic}/#")
    else:
        print(f"❌ MQTT连接失败，返回码 {rc}")

def on_message(client, userdata, msg):
    """MQTT消息回调"""
    try:
        # 处理设备状态消息
        if msg.topic == device_topic:
            payload = json.loads(msg.payload.decode())
            device_id = payload.get('device_id')
            
            if device_id:
                # 更新设备状态
                devices[device_id] = {
                    'name': payload.get('name', f'设备 {device_id}'),
                    'status': payload.get('status', 'unknown'),
                    'platform': payload.get('platform', 'unknown'),
                    'last_heartbeat': time.time()
                }
                print(f"🔄 设备 {device_id} 状态更新: {payload.get('status')}")
        
        # 处理图像feed消息
        elif msg.topic.startswith(feed_topic):
            handle_image_feed(msg)
            
    except Exception as e:
        print(f"❌ 处理MQTT消息时出错: {e}")

def handle_image_feed(msg):
    """处理图像feed消息"""
    try:
        payload = json.loads(msg.payload.decode())
        device_id = payload.get('device_id')
        image_type = payload.get('type')  # 'screen' 或 'camera'
        timestamp = payload.get('timestamp')
        
        if not device_id or not image_type or not timestamp:
            return
            
        # 解码图像
        image_data = base64.b64decode(payload.get('image', ''))
        if not image_data:
            return
            
        # 生成文件名和路径
        dt = datetime.fromisoformat(timestamp)
        date_str = dt.strftime("%Y-%m-%d")
        time_str = dt.strftime("%H-%M-%S")
        
        if image_type == 'screen':
            # 屏幕截图
            filename = f"{device_id}_screen_{date_str}_{time_str}.jpg"
            filepath = os.path.join("screenshots", filename)
        elif image_type == 'camera':
            # 摄像头截图
            filename = f"{device_id}_camera_{date_str}_{time_str}.jpg"
            filepath = os.path.join("camera_shots", filename)
        else:
            return
            
        # 保存图像
        with open(filepath, 'wb') as f:
            f.write(image_data)
            
        print(f"💾 已保存{image_type}截图: {filename}")
        
        # 根据图像类型进行相应处理
        if image_type == 'screen':
            # 屏幕截图 - 立即进行AI分析并保存结果到数据库
            analysis_result = ai_analyzer.analyze_screen_image(filepath, device_id, timestamp)
            if analysis_result:
                print(f"🧠 屏幕截图分析完成: {analysis_result[:100]}...")
                
        elif image_type == 'camera':
            # 摄像头截图 - 立即进行人体检测
            perform_human_detection(filepath, device_id, timestamp)
            
    except Exception as e:
        print(f"❌ 处理图像feed消息时出错: {e}")

def perform_human_detection(image_path, device_id, timestamp):
    """执行人体检测"""
    try:
        if pose_model is None:
            print("❌ 人体检测模型未加载")
            return
            
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            print(f"❌ 无法读取图像: {image_path}")
            return
            
        # 执行人体检测
        results = pose_model(image)
        
        # 检查是否检测到人体
        detected = False
        if len(results) > 0:
            for result in results:
                if result.keypoints is not None and len(result.keypoints.data) > 0:
                    detected = True
                    break
        
        # 处理检测结果
        if detected:
            # 检测到人体
            print(f"👤 检测到人体: {device_id} at {timestamp}")
            
            # 检查是否已有活跃的检测记录
            if device_id in active_detections:
                # 更新现有记录的结束时间
                log_id = active_detections[device_id]
                try:
                    detection_logger.end_detection(log_id, timestamp)
                    print(f"🕒 更新人体检测记录结束时间: {log_id}")
                except sqlite3.OperationalError as e:
                    if "database is locked" in str(e):
                        print(f"⚠️ 数据库被锁定，等待后重试: {e}")
                        time.sleep(1)  # 等待1秒后重试
                        try:
                            detection_logger.end_detection(log_id, timestamp)
                            print(f"🕒 重试成功，更新人体检测记录结束时间: {log_id}")
                        except Exception as retry_e:
                            print(f"❌ 重试更新人体检测记录失败: {retry_e}")
                    else:
                        print(f"❌ 更新人体检测记录失败: {e}")
                except Exception as e:
                    print(f"❌ 更新人体检测记录失败: {e}")
            else:
                # 创建新的人体检测记录
                try:
                    log_id = detection_logger.start_detection(device_id, timestamp)
                    active_detections[device_id] = log_id
                    print(f"⏺️ 开始新的人体检测记录: {log_id}")
                except sqlite3.OperationalError as e:
                    if "database is locked" in str(e):
                        print(f"⚠️ 数据库被锁定，等待后重试: {e}")
                        time.sleep(1)  # 等待1秒后重试
                        try:
                            log_id = detection_logger.start_detection(device_id, timestamp)
                            active_detections[device_id] = log_id
                            print(f"⏺️ 重试成功，开始新的人体检测记录: {log_id}")
                        except Exception as retry_e:
                            print(f"❌ 重试创建人体检测记录失败: {retry_e}")
                    else:
                        print(f"❌ 创建人体检测记录失败: {e}")
                except Exception as e:
                    print(f"❌ 创建人体检测记录失败: {e}")
        else:
            # 未检测到人体
            print(f"❌ 未检测到人体: {device_id} at {timestamp}")
            
            # 检查是否有活跃的检测记录需要结束
            if device_id in active_detections:
                # 结束记录并从活跃记录中移除
                log_id = active_detections[device_id]
                try:
                    detection_logger.end_detection(log_id, timestamp)
                    del active_detections[device_id]
                    print(f"⏹️ 结束人体检测记录: {log_id}")
                except sqlite3.OperationalError as e:
                    if "database is locked" in str(e):
                        print(f"⚠️ 数据库被锁定，等待后重试: {e}")
                        time.sleep(1)  # 等待1秒后重试
                        try:
                            detection_logger.end_detection(log_id, timestamp)
                            del active_detections[device_id]
                            print(f"⏹️ 重试成功，结束人体检测记录: {log_id}")
                        except Exception as retry_e:
                            print(f"❌ 重试结束人体检测记录失败: {retry_e}")
                    else:
                        print(f"❌ 结束人体检测记录失败: {e}")
                except Exception as e:
                    print(f"❌ 结束人体检测记录失败: {e}")
                
    except Exception as e:
        print(f"❌ 人体检测过程中出错: {e}")

# 设置MQTT客户端
client.username_pw_set(username, password)
client.on_connect = on_connect
client.on_message = on_message

# 连接MQTT服务器
try:
    client.connect(broker, port, 60)
    client.loop_start()
    print("📡 MQTT客户端已启动")
except Exception as e:
    print(f"❌ MQTT连接失败: {e}")

# Flask 路由
@app.route('/')
def index():
    """主页"""
    return render_template('index.html')

@app.route('/api/devices')
def get_devices():
    """获取当前设备列表（包括在线和离线设备）"""
    device_list = []
    for device_id, device_info in devices.items():
        device_list.append({
            "id": device_id,
            "name": device_info.get('name', f'设备 {device_id}'),
            "status": device_info.get('status', 'online'),  # 使用设备实际状态
            "last_heartbeat": device_info.get('last_heartbeat')
        })
    
    return jsonify(device_list)

@app.route('/api/history_devices')
def get_history_devices():
    """获取历史设备列表（包括在线和离线设备）"""
    try:
        # 获取当前所有设备（在线和离线）
        all_devices = {}
        for device_id, device_info in devices.items():
            all_devices[device_id] = {
                "id": device_id,
                "name": device_info.get('name', f'设备 {device_id}'),
                "status": device_info.get('status', 'online'),
                "last_heartbeat": device_info.get('last_heartbeat')
            }
        
        conn = get_db_connection()
        # 从数据库中获取所有曾经连接过的设备
        history_devices = conn.execute('''
            SELECT DISTINCT device_id FROM detection_logs
            UNION
            SELECT DISTINCT device_id FROM screen_analysis
            UNION
            SELECT DISTINCT device_id FROM daily_summary
            ORDER BY device_id
        ''').fetchall()
        conn.close()
        
        result = []
        for device in history_devices:
            device_id = device['device_id']
            # 检查设备是否在当前设备列表中
            if device_id in all_devices:
                result.append(all_devices[device_id])
            else:
                # 数据库中有记录但当前未连接的设备
                result.append({
                    "id": device_id,
                    "name": f'设备 {device_id}',
                    "status": "offline"
                })
        
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/device/<device_id>')
def get_device_detail(device_id):
    """获取设备详情"""
    try:
        # 获取设备基本信息
        device_info = devices.get(device_id, {})
        
        # 获取设备的人体检测记录
        conn = get_db_connection()
        
        # 获取最近的人体检测记录
        detection_logs = conn.execute(
            'SELECT * FROM detection_logs WHERE device_id = ? ORDER BY start_time DESC LIMIT 10',
            (device_id,)
        ).fetchall()
        
        # 获取设备的屏幕分析记录（包含Base64编码的图片）
        screen_analyses = conn.execute(
            'SELECT * FROM screen_analysis WHERE device_id = ? ORDER BY timestamp DESC LIMIT 5',
            (device_id,)
        ).fetchall()
        
        # 获取设备的每日统计
        daily_stats = conn.execute(
            'SELECT * FROM daily_summary WHERE device_id = ? ORDER BY date DESC LIMIT 30',
            (device_id,)
        ).fetchall()
        
        conn.close()
        
        # 格式化数据
        detection_data = []
        for log in detection_logs:
            detection_data.append({
                "id": log["id"],
                "date": log["date"],
                "start_time": log["start_time"],
                "end_time": log["end_time"],
                "duration": log["duration"]
            })
        
        screen_data = []
        for analysis in screen_analyses:
            # 将图片编码为Base64
            base64_image = encode_image_to_base64(analysis["image_path"])
            screen_data.append({
                "id": analysis["id"],
                "date": analysis["date"],
                "timestamp": analysis["timestamp"],
                "image_path": analysis["image_path"],
                "analysis_result": analysis["analysis_result"],
                "base64_image": base64_image if base64_image else ""
            })
        
        stats_data = []
        for stat in daily_stats:
            stats_data.append({
                "id": stat["id"],
                "date": stat["date"],
                "total_duration": stat["total_duration"],
                "detection_count": stat["detection_count"]
            })
        
        result = {
            "device_id": device_id,
            "device_info": device_info,
            "detection_logs": detection_data,
            "screen_analyses": screen_data,
            "daily_stats": stats_data
        }
        
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/device/<device_id>')
def device_detail_page(device_id):
    """设备详情页面"""
    return render_template('device_detail.html', device_id=device_id)

@app.route('/api/device_stats')
def get_device_stats():
    """获取设备统计信息"""
    online_count = len(devices)
    # 在实际项目中，你可能需要跟踪总设备数
    total_count = online_count  # 简化处理
    
    stats = {
        "online": online_count,
        "total": total_count
    }
    return jsonify(stats)

@app.route('/api/daily_summary')
def get_daily_summary():
    """获取每日总结"""
    try:
        conn = get_db_connection()
        # 获取最近7天的总结
        seven_days_ago = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d")
        summaries = conn.execute(
            'SELECT * FROM daily_activities WHERE date >= ? ORDER BY date DESC',
            (seven_days_ago,)
        ).fetchall()
        conn.close()
        
        result = []
        for summary in summaries:
            result.append({
                "id": summary["id"],
                "device_id": summary["device_id"],
                "date": summary["date"],
                "summary": summary["summary"],
                "detailed_analysis": summary["detailed_analysis"],
                "created_at": summary["created_at"]
            })
        
        return jsonify(result)
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/human_detection_stats')
def get_human_detection_stats():
    """获取人体检测统计信息"""
    try:
        conn = get_db_connection()
        # 获取最近30天的统计数据
        thirty_days_ago = (datetime.now() - timedelta(days=30)).strftime("%Y-%m-%d")
        stats = conn.execute(
            'SELECT date, total_duration, detection_count FROM daily_summary WHERE date >= ? ORDER BY date',
            (thirty_days_ago,)
        ).fetchall()
        conn.close()
        
        dates = []
        durations = []
        counts = []
        
        for stat in stats:
            dates.append(stat["date"])
            durations.append(stat["total_duration"])
            counts.append(stat["detection_count"])
        
        return jsonify({
            "dates": dates,
            "durations": durations,
            "counts": counts
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/screen_analysis_stats')
def get_screen_analysis_stats():
    """获取屏幕分析统计信息"""
    try:
        conn = get_db_connection()
        # 获取最近30天的屏幕分析统计数据
        thirty_days_ago = (datetime.now() - timedelta(days=30)).strftime("%Y-%m-%d")
        stats = conn.execute(
            'SELECT date, COUNT(*) as analysis_count FROM screen_analysis WHERE date >= ? GROUP BY date ORDER BY date',
            (thirty_days_ago,)
        ).fetchall()
        conn.close()
        
        dates = []
        counts = []
        
        for stat in stats:
            dates.append(stat["date"])
            counts.append(stat["analysis_count"])
        
        return jsonify({
            "dates": dates,
            "counts": counts
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/detection_distribution')
def get_detection_distribution():
    """获取人体检测时间分布"""
    try:
        conn = get_db_connection()
        # 获取最近7天的人体检测时间分布（按小时统计）
        seven_days_ago = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d")
        distributions = conn.execute(
            '''SELECT 
                strftime('%H', start_time) as hour,
                COUNT(*) as count
            FROM detection_logs 
            WHERE date >= ? 
            GROUP BY hour 
            ORDER BY hour''',
            (seven_days_ago,)
        ).fetchall()
        conn.close()
        
        hours = []
        counts = []
        
        for dist in distributions:
            hours.append(dist["hour"] + ":00")
            counts.append(dist["count"])
        
        return jsonify({
            "hours": hours,
            "counts": counts
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/api/screenshots')
def get_screenshots():
    """获取截图列表"""
    try:
        # 限制截图文件夹中的文件数量
        limit_image_files("screenshots")
        limit_image_files("camera_shots")
        limit_image_files("detections")
        
        # 获取所有截图文件列表
        all_files = []
        for folder in ["screenshots", "camera_shots", "detections"]:
            files = glob.glob(f"{folder}/*.jpg")
            all_files.extend(files)
        
        # 按修改时间排序，获取最新的3张图片
        all_files.sort(key=os.path.getmtime, reverse=True)
        latest_files = all_files[:3]
        
        # 处理最新的3张图片
        screenshots = []
        for file in latest_files:
            # 将图片编码为Base64
            base64_image = encode_image_to_base64(file)
            if base64_image:
                screenshots.append({
                    "path": file.replace("\\", "/"),
                    "name": os.path.basename(file),
                    "folder": file.split("/")[0],
                    "mtime": os.path.getmtime(file),
                    "base64": base64_image
                })
        
        return jsonify(screenshots)  # 返回最多3张最新图片
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    # 启动Flask应用
    app.run(host=FLASK_CONFIG['host'], 
            port=FLASK_CONFIG['port'], 
            debug=FLASK_CONFIG['debug'], 
            use_reloader=False)